package sm4

import (
	"encoding/hex"
	"fmt"

	"github.com/tjfoc/gmsm/sm4"
)

type SM4ECB struct {
	key []byte
}

func NewSM4ECB(key []byte) *SM4ECB {
	return &SM4ECB{
		key: key,
	}
}

// SM4ECBEncrypt SM4/ECB/PKCS5Padding 加密
func (a *SM4ECB) SM4ECBEncrypt(plainText, key []byte) ([]byte, error) {
	// 验证密钥长度
	if len(key) != sm4.BlockSize {
		return nil, fmt.Errorf("invalid key size %d, expect %d", len(key), sm4.BlockSize)
	}
	// 创建 SM4 密码块
	block, err := sm4.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("failed to create SM4 cipher: %v", err)
	}
	// PKCS5Padding 填充
	plainText = a.PKCS5Padding(plainText, block.BlockSize())
	// ECB 模式加密
	cipherText := make([]byte, len(plainText))
	blockSize := block.BlockSize()
	for i := 0; i < len(plainText); i += blockSize {
		block.Encrypt(cipherText[i:i+blockSize], plainText[i:i+blockSize])
	}
	return cipherText, nil
}

// SM4ECBDecrypt SM4/ECB/PKCS5Padding 解密
func (a *SM4ECB) SM4ECBDecrypt(cipherText, key []byte) ([]byte, error) {
	// 验证密钥长度
	if len(key) != sm4.BlockSize {
		return nil, fmt.Errorf("invalid key size %d, expect %d", len(key), sm4.BlockSize)
	}
	// 验证密文长度
	if len(cipherText)%sm4.BlockSize != 0 {
		return nil, fmt.Errorf("invalid ciphertext size %d", len(cipherText))
	}
	// 创建 SM4 密码块
	block, err := sm4.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("failed to create SM4 cipher: %v", err)
	}
	// ECB 模式解密
	plainText := make([]byte, len(cipherText))
	blockSize := block.BlockSize()
	for i := 0; i < len(cipherText); i += blockSize {
		block.Decrypt(plainText[i:i+blockSize], cipherText[i:i+blockSize])
	}
	// PKCS5Unpadding 去除填充
	plainText, err = a.PKCS5Unpadding(plainText)
	if err != nil {
		return nil, fmt.Errorf("failed to unpadding: %v", err)
	}
	return plainText, nil
}

// PKCS5Padding PKCS5 填充
func (a *SM4ECB) PKCS5Padding(src []byte, blockSize int) []byte {
	padding := blockSize - len(src)%blockSize
	padtext := make([]byte, padding)
	for i := range padtext {
		padtext[i] = byte(padding)
	}
	return append(src, padtext...)
}

// PKCS5Unpadding PKCS5 去除填充
func (a *SM4ECB) PKCS5Unpadding(src []byte) ([]byte, error) {
	length := len(src)
	if length == 0 {
		return nil, fmt.Errorf("invalid padding length")
	}
	unpadding := int(src[length-1])
	if unpadding > length {
		return nil, fmt.Errorf("invalid padding size")
	}
	return src[:(length - unpadding)], nil
}

func (a *SM4ECB) Encrypt(plainText []byte) (string, error) {
	cipherText, err := a.SM4ECBEncrypt(plainText, a.key)
	if err != nil {
		return "", err
	}
	return hex.EncodeToString(cipherText), nil
}

func (a *SM4ECB) Decrypt(cipherText string) (string, error) {
	cipherTextBytes, err := hex.DecodeString(cipherText)
	if err != nil {
		return "", err
	}
	decrypted, err := a.SM4ECBDecrypt(cipherTextBytes, a.key)
	if err != nil {
		return "", err
	}
	return string(decrypted), nil
}
